import numpy as np
from problems.TSP import initialization, fitness, crossover, mutation, selection, generate, draw
from problems.TSP.problem_instance import TSPInstance
from local_search_operators.apply_local_search import apply_local_search


# 主函数 - 遗传算法求解TSP
def genetic_algorithm_tsp(seed, num_cities, pop_size, elite_size, mutation_rate, generations, ls_intensity,operator_names):
    # 生成城市坐标
    cities = initialization.generate_cities(num_cities, seed)

    # 创建TSP问题实例
    problem_instance = TSPInstance(cities)

    # 1. 初始化种群
    population = initialization.initialize_population(pop_size, cities)

    # 记录最佳解和进化历史
    best_individual = None
    best_fitness = -np.inf
    history = []

    # 迭代进化
    for gen in range(generations):
        # 2. 计算适应度
        fitness_values = fitness.calculate_fitness(population, cities)

        # 3. 记录当前最佳解
        current_best_idx = np.argmax(fitness_values)
        current_best_fitness = fitness_values[current_best_idx]
        current_best_individual = population[current_best_idx]

        if current_best_fitness > best_fitness:
            best_fitness = current_best_fitness
            best_individual = current_best_individual.copy()

        # 记录历史数据
        history.append(1 / best_fitness)

        # 4. 选择操作
        selected_population = selection.roulette_wheel_selection(population, fitness_values)

        # 5. 交叉操作
        new_population = []
        for i in range(0, pop_size, 2):
            parent1 = selected_population[i]
            parent2 = selected_population[i + 1]

            # 交叉产生两个后代
            child1 = crossover.order_crossover(parent1, parent2)
            child2 = crossover.order_crossover(parent2, parent1)

            new_population.extend([child1, child2])

        # 6. 变异操作
        for i in range(len(new_population)):
            new_population[i] = mutation.swap_mutation(new_population[i], mutation_rate)

        # 7. 精英保留策略
        new_population = generate.elitism(population, fitness_values, new_population, elite_size)

        # 8. 应用局部搜索
        if ls_intensity > 0:
            new_population = apply_local_search(
                new_population,
                problem_type='tsp',
                problem_instance=problem_instance,
                elite_size=elite_size,
                ls_intensity=ls_intensity,
                operator_names=operator_names
            )

        # 更新种群
        population = np.array(new_population)

        # 每50代打印进度
        if gen % 50 == 0:
            print(f"Generation {gen}: Best Distance = {1 / best_fitness:.2f}")

    # 最终结果可视化
    best_distance = 1 / best_fitness
    print("\nGenetic Algorithm Completed!")
    print(f"Best Distance: {best_distance:.2f}")
    print(f"Best Path: {best_individual}")

    # 绘制结果
    draw.plot_results(history)

    # 绘制最佳路径
    draw.plot_path(cities, best_individual, f"Best TSP Path (Distance: {best_distance:.2f})")

    return best_individual, best_distance